from transformers import AutoTokenizer , AutoModelForSequenceClassification
import torch

tokenizer = AutoTokenizer.from_pretrained('/home/models/bert-base')
model = AutoModelForSequenceClassification.from_pretrained('/home/models/minilm')

def rerank(query,documents):
    inputs = tokenizer([f"{query} {doc}" for doc in documents],padding=True,return_tensors='pt')
    with torch.no_grad():
        outputs = model(**inputs)
        scroes = torch.softmax(outputs.logits,dim=-1)[:-1].tolist()
        ranked_docs = sorted(zip(scroes,documents),key=lambda x:x[0],reverse=True)
        return [(score,doc) for score,doc in ranked_docs]

query = "what is the capital of France?"
documents = [
    "Parais is the catital city of France",
    "what is the capital of France?"
]
ranked_results = rerank(query,documents)
print(ranked_results)